# 盲源信号分离实战
import numpy as np
import matplotlib.pyplot as plt
from scipy import signal #使用scipy库
from sklearn.decomposition import FastICA

# 数据生成
np.random.seed(0)
n_samples = 200
time = np.linspace(0,8,n_samples) # 生成采样数据
s1 = np.sin(2 * time) # 信号1：正弦信号
s2 = np.sign(np.sin(3 * time)) # 信号2：方波信号
s3 = signal.sawtooth(2*np.pi*time) # 信号3 锯齿波信号
S = np.c_[s1,s2,s3] # 将对象沿第二个轴（按列）连接，shape为（200,3）
S += 0.1 * np.random.normal(size=S.shape) # 增加噪声
S /= S.std(axis=0) # 数据归一化
A = np.array([[1,1,1],[0.5,2,1.0],[1.5,1.0,2.0],[1.0,5.0,0.75]]) # 混淆矩阵
Y = np.dot(S,A.T) # 生成观测信号，shape为（200,4）每列代表一个观测信号

ica = FastICA(n_components=3)
S_ = ica.fit_transform(Y) # 重构原始信号，shape为(200,3)
A_ = ica.mixing_ #获得估计混淆矩阵，shape为(4,3)

# 开始绘图
plt.figure()
models = [Y, S, S_, ]
names = ["Observation signals","True Sources","ICA recovered signals"]
colors = ["red","steelblue","orange","blue"]
for ii, (model,name) in enumerate(zip(models,names),1):
      plt.subplot(3,1,ii) # 准备画每个子图
      plt.title(name)
      for sig,color in zip(model.T, colors):
            plt.plot(sig,color=color)
plt.tight_layout() # 自动调整子图大小
plt.show()